import matplotlib.pyplot as plt
import pandas as pd
import os
import argparse
import shutil
import numpy as np

def plotAirfoilGeom(coordinates, file_path=None):
    '''
        绘制翼型
    :param coordinates:
    :param file_path:
    :return:
    '''
    len = int(coordinates.shape[0] / 2)

    plt.scatter(coordinates[0:len, 0], coordinates[0:len, 1],
                color='red',
                s=1)
    plt.scatter(coordinates[len::, 0], coordinates[len::, 1],
                color='blue',
                s=1)

    # plt.plot(coordinates[0:len, 0], coordinates[0:len, 1])
    # plt.plot(coordinates[len::, 0], coordinates[len::, 1])

    plt.xlim((-0.01, 1.01))
    plt.ylim((-0.3, 0.3))
    # plt.xlim((-0.01, 0.15))
    # plt.ylim((-0.08, 0.08))
    plt.gca().set_aspect(1)
    plt.grid()
    plt.show()
    # plt.savefig(file_path, dpi=800)
    plt.close()


def plotAirfoilGeom2(coordinates, file_name, save_file_path=None):
    color_list = ["red", "blue", "black"]

    for i, crd in enumerate(coordinates):
        len = int(crd.shape[0] / 2)
        # plt.scatter(crd[0:len, 0], crd[0:len, 1],
        #             color=color_list[i],
        #             s=10)  # , linestyle='-')
        # plt.scatter(crd[len::, 0], crd[len::, 1],
        #             color=color_list[i],
        #             s=10)  # , linestyle='-')
        plt.plot(crd[0:len + 1, 0], crd[0:len + 1, 1],
                 color=color_list[i])
        plt.plot(crd[len::, 0], crd[len::, 1],
                 color=color_list[i],
                 label=file_name[i][:-4])

    for i in range(coordinates.shape[0] - 1):
        e = coordinates[i + 1]
        s = coordinates[i]
        plt.annotate("",
                     xy=(e[0], e[1]),
                     xytext=(s[0], s[1]),
                     arrowprops=dict(arrowstyle="->", color="r"))

    plt.xlim((-0.02, 1.02))
    plt.ylim((-0.3, 0.3))
    plt.gca().set_aspect(1)
    plt.grid()
    plt.legend()
    plt.savefig(save_file_path, dpi=600, bbox_inches='tight')
    plt.close()


def plotAirfoilGeom3(coordinates_list, file_path=None):
    '''
        绘制翼型
    :param coordinates:
    :param file_path:
    :return:
    '''

    labels = ["cst", "bspline"]
    ld = [2.0, 1.0]
    ls = [':', '--']

    for i, coordinates in enumerate(coordinates_list):
        len = int(coordinates.shape[0] / 2)
        plt.plot(coordinates[0:len, 0], coordinates[0:len, 1]*3, linewidth=ld[i],
                 linestyle=ls[i], label=labels[i])
        plt.plot(coordinates[len::, 0], coordinates[len::, 1]*3, linewidth=ld[i],
                 linestyle=ls[i], label=labels[i])

    plt.xlim((-0.01, 1.01))
    plt.ylim((-0.3, 0.3))
    # plt.xlim((-0.01, 0.05))
    # plt.ylim((-0.3, 0.3))
    # plt.legend()
    plt.gca().set_aspect(1)
    plt.grid()
    plt.savefig(file_path, dpi=1000)
    plt.close()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-r", "--read_dir", type=str, help="待处理翼型所在文件夹路径", required=True)
    parser.add_argument("-s", "--save_dir", type=str, help="处理后的翼型保存路径", required=True)
    args = parser.parse_args()

    if not os.path.exists(args.save_dir):
        os.mkdir(args.save_dir)

    # 画单个翼型
    # for file in os.listdir(args.read_dir):
    #     file_path = os.path.join(args.read_dir, file)
    #     save_path = os.path.join(args.save_dir, file[:-3] + "jpg")
    #     coordinates1 = np.loadtxt(file_path, skiprows=1)
    #     plotAirfoilGeom(coordinates1, save_path)
    # print("二维翼型可视化图已保存在{}路径下".format(args.save_dir))

    # 画多个翼型
    for file in os.listdir(args.read_dir):
        file_path1 = os.path.join(args.read_dir, file)
        file = "rae2822.csv"
        file_path2 = os.path.join(r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_11", file)

        save_path = os.path.join(args.save_dir, file[:-3] + "jpg")
        coordinates1 = np.loadtxt(file_path1, skiprows=1)
        coordinates2 = pd.read_csv(file_path2, header=0, index_col=None).to_numpy()
        plotAirfoilGeom3([coordinates1, coordinates2], save_path)
    print("二维翼型可视化图已保存在{}路径下".format(args.save_dir))
